import concurrent.futures
import time
from datetime import datetime, timedelta
from typing import List, Optional

import akshare as ak
import pandas as pd
import sqlalchemy as sa
from sqlalchemy.dialects.postgresql import insert
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from tqdm import tqdm

from .database_config import DatabaseConfig

Base = declarative_base()


class StockKline(Base):
    __tablename__ = 'stock_kline'

    ts_code = sa.Column(sa.CHAR(9), primary_key=True)
    trade_date = sa.Column(sa.DateTime(timezone=True), primary_key=True)
    open = sa.Column(sa.Numeric(8, 2))
    high = sa.Column(sa.Numeric(8, 2))
    low = sa.Column(sa.Numeric(8, 2))
    close = sa.Column(sa.Numeric(8, 2))
    volume = sa.Column(sa.BigInteger)


class MarketDataFetcher:
    def __init__(self, db_config: Optional[DatabaseConfig] = None, max_workers: int = 20,
                 retry_count: int = 3, delay: float = 0.1):
        """
        Initialize MarketDataFetcher with PostgreSQL connection
        
        Args:
            db_config: Database configuration
            max_workers: Maximum number of worker threads
            retry_count: Number of retries for failed requests
            delay: Delay between retries in seconds
        """
        self.db_config = db_config or DatabaseConfig()
        self.max_workers = max_workers
        self.retry_count = retry_count
        self.delay = delay

        self.engine = sa.create_engine(self.db_config.connection_string)
        Base.metadata.create_all(self.engine)
        self.Session = sessionmaker(bind=self.engine)

    def _clean_data(self, df: pd.DataFrame) -> pd.DataFrame:
        """Clean and validate market data"""
        if df.empty:
            return df

        # Convert date column, yyyy-MM-dd

        df['trade_date'] = pd.to_datetime(df['日期']).dt.strftime('%Y-%m-%d')
        # Rename columns to match database schema
        column_mapping = {
            '开盘': 'open',
            '收盘': 'close',
            '最高': 'high',
            '最低': 'low',
            '成交量': 'volume'
        }
        df = df.rename(columns=column_mapping)

        # Select and reorder columns
        columns = ['trade_date', 'open', 'high', 'low', 'close', 'volume']
        df = df[columns]

        # Remove rows with missing values
        df = df.dropna()

        # Remove extreme values (e.g., prices <= 0)
        df = df[df['close'] > 0]

        # Convert numeric columns to appropriate types
        for col in ['open', 'high', 'low', 'close']:
            df[col] = pd.to_numeric(df[col], errors='coerce').round(2)
        df['volume'] = pd.to_numeric(df['volume'], errors='coerce').astype('int64')

        return df

    def get_stock_codes(self) -> List[str]:
        """Get list of all A-stock codes"""
        try:
            stock_info_df = ak.stock_info_a_code_name()
            return [code for code in stock_info_df['code'].tolist()]
        except Exception as e:
            print(f"Failed to get stock codes: {str(e)}")
            return []

    def _get_latest_date(self, ts_code: str) -> Optional[datetime]:
        """Get the latest date for a symbol from database"""
        with self.Session() as session:
            result = session.query(sa.func.max(StockKline.trade_date)) \
                .filter(StockKline.ts_code == ts_code) \
                .scalar()
        return result

    def _fetch_stock_data(self, ts_code: str, end_date: datetime,
                          start_date: Optional[datetime] = None) -> Optional[pd.DataFrame]:
        """Fetch stock data for a given symbol and date range, fetching backwards from end_date in 500-day segments"""
        latest_date = self._get_latest_date(ts_code)
        symbol = ts_code.split('.')[0]  # Remove exchange suffix for akshare

        if latest_date is not None:
            if end_date <= latest_date:
                print(f"{ts_code}: Data is up to date")
                return None
            start_date = latest_date + timedelta(days=1)
        elif start_date is None:
            start_date = datetime(1980, 12, 19)

        print(f"{ts_code}: Fetching data from {start_date.date()} to {end_date.date()}")

        all_data = []
        current_end = end_date
        segment_size = timedelta(days=1000)

        while current_end >= start_date:
            current_start = max(current_end - segment_size, start_date)
            
            for attempt in range(self.retry_count):
                try:
                    df = ak.stock_zh_a_hist(symbol=symbol,
                                          start_date=current_start.strftime('%Y%m%d'),
                                          end_date=current_end.strftime('%Y%m%d'),
                                          adjust='qfq')
                    if df is not None and not df.empty:
                        df = self._clean_data(df)
                        if df is not None and not df.empty:
                            df['ts_code'] = ts_code
                            all_data.append(df)
                            break
                    time.sleep(self.delay)
                except Exception as e:
                    print(f"Attempt {attempt + 1} failed for {ts_code} from {current_start.date()} to {current_end.date()}: {str(e)}")
                    if attempt < self.retry_count - 1:  # 如果不是最后一次重试，则等待后重试
                        time.sleep(self.delay * (attempt + 1))
                    else:  # 如果是最后一次重试失败，打印警告并继续下一个时间段
                        print(f"Warning: Failed to fetch data for period {current_start.date()} to {current_end.date()}")
            
            current_end = current_start - timedelta(days=1)

        if all_data:
            # 合并数据并按日期排序
            result = pd.concat(all_data, ignore_index=True)
            return result.sort_values('trade_date').reset_index(drop=True)
        return None

    def _save_to_db(self, ts_code: str, df: pd.DataFrame) -> None:
        """Save data to PostgreSQL database"""
        if df is None or df.empty:
            return

        records = df.to_dict('records')

        with self.Session() as session:
            try:
                stmt = insert(StockKline).values(records)
                stmt = stmt.on_conflict_do_nothing(
                    index_elements=['ts_code', 'trade_date']
                )
                session.execute(stmt)
                session.commit()
            except Exception as e:
                print(f"Failed to save data for {ts_code}: {str(e)}")
                session.rollback()

    def update_stock_data(self, symbols: Optional[List[str]] = None,
                          end_date: Optional[datetime] = None) -> None:
        """Update stock data for given symbols"""
        if symbols is None:
            symbols = self.get_stock_codes()

        if end_date is None:
            end_date = datetime.now()

        with concurrent.futures.ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            futures = []
            for symbol in symbols:
                future = executor.submit(self._fetch_stock_data, symbol, end_date)
                futures.append((symbol, future))

            for symbol, future in tqdm(futures, desc="Updating stock data"):
                try:
                    df = future.result()
                    if df is not None:
                        self._save_to_db(symbol, df)
                except Exception as e:
                    print(f"Failed to process {symbol}: {str(e)}")



    def get_stock_data(self, ts_code: str, start_date: Optional[datetime] = None,
                       end_date: Optional[datetime] = None) -> pd.DataFrame:
        """Retrieve stock data from database"""
        query = sa.select(StockKline).filter(StockKline.ts_code == ts_code)

        if start_date:
            query = query.filter(StockKline.trade_date >= start_date)
        if end_date:
            query = query.filter(StockKline.trade_date <= end_date)

        query = query.order_by(StockKline.trade_date)

        with self.Session() as session:
            results = session.execute(query).all()
            records = []
            for row in results:
                stock = row[0]  # 获取StockKline对象
                record = {
                    'ts_code': stock.ts_code,
                    'trade_date': stock.trade_date,
                    'open': float(stock.open) if stock.open else None,
                    'high': float(stock.high) if stock.high else None,
                    'low': float(stock.low) if stock.low else None,
                    'close': float(stock.close) if stock.close else None,
                    'volume': int(stock.volume) if stock.volume else None
                }
                records.append(record)

            df = pd.DataFrame.from_records(records)

            # 确保按日期排序
            if not df.empty:
                df = df.sort_values('trade_date').reset_index(drop=True)

        return df

    def sync_realtime_market_data(self) -> pd.DataFrame:

        result = self.get_realtime_market_data()

        # 只保留需要的列
        result = result[['ts_code', 'open', 'high', 'low', 'close', 'volume']]

        try:
            # 保存到数据库
            records = result.to_dict('records')
            
            with self.Session() as session:
                stmt = insert(StockKline).values(records)
                stmt = stmt.on_conflict_do_nothing(
                    index_elements=['ts_code', 'trade_date']
                )
                session.execute(stmt)
                session.commit()
                
            print(f"Successfully updated {len(result)} records to StockKline table")
            
        except Exception as e:
            print(f"Error saving data to StockKline table: {str(e)}")
        
        return result

    def get_realtime_market_data(self) -> pd.DataFrame:
        """
                同步所有市场的实时行情数据到 StockKline 表
                包括：深圳、上海、北京、创业板、科创板等市场，
                每天行情结束后
                """
        all_market_data = []

        # 定义所有需要获取的市场及其对应的函数
        market_functions = {
            'SZ': ak.stock_sz_a_spot_em,  # 深圳A股
            'BJ': ak.stock_bj_a_spot_em,  # 北京A股
            'CY': ak.stock_cy_a_spot_em,  # 创业板
            'KC': ak.stock_kc_a_spot_em,  # 科创板
            'NEW': ak.stock_new_a_spot_em  # 新股
        }

        current_time = pd.Timestamp.now().strftime('%Y-%m-%d')

        for market, func in market_functions.items():
            try:
                print(f"Fetching {market} market data...")
                df = func()
                if df is not None and not df.empty:
                    # 统一字段名称
                    column_mapping = {
                        '代码': 'ts_code',  # 股票代码
                        '名称': 'name',  # 股票名称
                        '最新价': 'close',  # 收盘价
                        '涨跌幅': 'pct_change',  # 涨跌幅（百分比变化）
                        '涨跌额': 'change',  # 涨跌额
                        '成交量': 'volume',  # 成交量
                        '成交额': 'turnover',  # 成交金额
                        '振幅': 'amplitude',  # 振幅
                        '最高': 'high',  # 最高价
                        '最低': 'low',  # 最低价
                        '今开': 'open',  # 今日开盘价
                        '昨收': 'prev_close',  # 昨日收盘价
                        '量比': 'volume_ratio',  # 量比
                        '换手率': 'turnover_rate',  # 换手率
                        '市盈率-动态': 'pe_dynamic',  # 动态市盈率
                        '市净率': 'pb_ratio',  # 市净率
                        '总市值': 'total_market_cap',  # 总市值
                        '流通市值': 'circulating_market_cap',  # 流通市值
                        '涨速': 'price_change_rate',  # 涨速
                        '5分钟涨跌': 'five_min_change',  # 5分钟涨跌
                        '60日涨跌幅': 'sixty_day_change',  # 60日涨跌幅
                        '年初至今涨跌幅': 'ytd_change'  # 年初至今涨跌幅
                    }

                    # 重命名列
                    df = df.rename(columns=lambda x: column_mapping.get(x, x))

                    df.dropna(how='any', inplace=True)

                    # 将数值类型的列转换为数值
                    numeric_columns = ['open', 'high', 'low', 'close']
                    for col in numeric_columns:
                        if col in df.columns:
                            df[col] = pd.to_numeric(df[col], errors='coerce').round(2)

                    df['volume'] = pd.to_numeric(df['volume'], errors='coerce').astype('int64')

                    # 添加交易日期
                    df['trade_date'] = current_time

                    all_market_data.append(df)
                    print(f"Successfully fetched {market} market data: {len(df)} records")

                time.sleep(self.delay)  # 避免请求过快

            except Exception as e:
                print(f"Error fetching {market} market data: {str(e)}")
                continue

        if not all_market_data:
            return pd.DataFrame()

        # 合并所有市场数据
        result = pd.concat(all_market_data, ignore_index=True)
        return result
